/*
 * This is written by Lulu Han.
 * E-mail: locomotive_crypto@163.com
 */

#include "paillier.h"

Paillier::PublicKey::PublicKey(){
	// ...
}

Paillier::PublicKey::PublicKey(const NTL::ZZ& N){
	m_N = N;
	m_g = N + 1;
	m_N_square = N * N;
}

const NTL::ZZ Paillier::PublicKey::GetN(){
	return m_N;
}

const NTL::ZZ Paillier::PublicKey::GetG(){
	return m_g;
}

const NTL::ZZ Paillier::PublicKey::GetNSquare(){
	return m_N_square;
}

Paillier::PrivateKey::PrivateKey():PublicKey(){
	// ...
}

Paillier::PrivateKey::PrivateKey(const NTL::ZZ& p, const NTL::ZZ& q):PublicKey(p*q){
	m_Phi_N = (p-1)*(q-1);
	m_Phi_N_inv = NTL::InvMod(m_Phi_N, m_N);
}

const NTL::ZZ Paillier::PrivateKey::GetPhiN(){
	return m_Phi_N;
}

const NTL::ZZ Paillier::PrivateKey::GetPhiNInv(){
	return m_Phi_N_inv;
}
		
Paillier::Encryptor::Encryptor(const PublicKey& pub):PublicKey(pub){
	// ...
}

void Paillier::Encryptor::Encrypt(const NTL::ZZ& m, NTL::ZZ& c){
	NTL::ZZ r; //  1 < r < N
	while(true){
		r = NTL::RandomBnd(GetN());
		if (NTL::GCD(r, GetN()) == 1 && r > 1)
			break;
	}
	
	// Encrypt m
	NTL::ZZ tmp;
	tmp = NTL::PowerMod(r, GetN(), GetNSquare()); 
	NTL::mul(c, m, GetN());
	c = NTL::AddMod(c, 1, GetNSquare());
	c = NTL::MulMod(tmp, c, GetNSquare());
}

void Paillier::Encryptor::HE_Add_Cipher(const NTL::ZZ& c1, const NTL::ZZ& c2, NTL::ZZ& c){
	c = NTL::MulMod(c1, c2, GetNSquare());
}

void Paillier::Encryptor::HE_Add_Const(const NTL::ZZ& c1, const NTL::ZZ& n, NTL::ZZ& c){
	NTL::ZZ tmp;
	this->Encrypt(n, tmp);
	this->HE_Add_Cipher(c1, tmp, c);
}

void Paillier::Encryptor::HE_Mul_Const(const NTL::ZZ& c1, const NTL::ZZ& n, NTL::ZZ& c){
	c = NTL::PowerMod(c1, n, GetNSquare());
}
		
Paillier::Decryptor::Decryptor(const PrivateKey& pri):PrivateKey(pri){
	// ...
}


void Paillier::Decryptor::Decrypt(const NTL::ZZ& c, NTL::ZZ& m){
	NTL::ZZ tmp;
	tmp = NTL::PowerMod(c, GetPhiN(), GetNSquare());
	tmp -= 1;
	NTL::divide(tmp, tmp, GetN());
	m = NTL::MulMod(tmp, GetPhiNInv(), GetN());
}


void Paillier::Decryptor::Encrypt(const NTL::ZZ& m, NTL::ZZ& c){
	NTL::ZZ r; //  1 < r < N
	while(true){
		r = NTL::RandomBnd(GetN());
		if (NTL::GCD(r, GetN()) == 1 && r > 1)
			break;
	}
	
	// Encrypt m
	NTL::ZZ tmp;
	tmp = NTL::PowerMod(r, GetN(), GetNSquare()); 
	NTL::mul(c, m, GetN());
	c = NTL::AddMod(c, 1, GetNSquare());
	c = NTL::MulMod(tmp, c, GetNSquare());
}

void Paillier::GenKeyPair(PrivateKey& pri, PublicKey& pub, long bitlens){
	
	NTL::ZZ p, q;
	
	// Generate random prime number p and q that they are all bitlens length
	NTL::RandomPrime(p, bitlens, 80);
	NTL::RandomPrime(q, bitlens, 80);
	
	// Construct Private key and public key objects using p and q generated previously
	PrivateKey prikey(p, q);
	PublicKey pubkey(p*q);
	
	pri = prikey;
	pub = pubkey;
}

void Paillier::GenKeyPair(PrivateKey& pri, PublicKey& pub, const NTL::ZZ& p, const NTL::ZZ& q){
	// Construct Private key and public key objects according to user-input p and q
	// Assume that p and q is prime numbers
	PrivateKey prikey(p, q);
	PublicKey pubkey(p*q);
	
	pri = prikey;
	pub = pubkey;
}

